open Llvm
open Llvm_scalar_opts

exception Error of string

let (@) = fun f g x -> f(g(x))

let context = global_context ()
let lmodule = create_module context "mainModule"
let builder = builder context

(* Basic types *)
let void_type = void_type context
let boolean_type = integer_type context 1
let i8_type = i8_type context
let i32_type = i32_type context
let i64_type = i64_type context
let i8p_type = pointer_type i8_type
let i8pp_type = pointer_type i8p_type  

(* Use native integer size *)
let integer_type = match Sys.word_size with
  | 32 -> i32_type
  | 64 -> i64_type
  | _ -> failwith "Unrecognized word size"

let rec codegen_ltype = function
  | Ast.Boolean -> boolean_type
  | Ast.Integer -> integer_type
  | Ast.Void -> void_type
  | Ast.Record name -> let rtype = List.map snd (Types.fields_sort name) in
    (* Replace Record with Pointer so we can use recursive definitions *)
    let rtype = List.map (function Ast.Record _ -> Ast.Pointer | mptype -> mptype) rtype in
      pointer_type (struct_type context (Array.of_list (List.map codegen_ltype rtype)))
  | Ast.Pointer -> i8p_type;
;;

let get_function fname = match lookup_function fname lmodule with
  | Some lfunc -> lfunc
  | None -> failwith ("function not found " ^ fname)
;;

let codegen_call fname largs =
  (* Use empty name, because we can't assing one to calls returning void *)
  build_call (get_function fname) largs "" builder
;;

(* Declare void @llvm.gcroot( i8** alloca, i8* metedata ) *)
declare_function "llvm.gcroot" (function_type void_type [| i8pp_type; i8p_type |]) lmodule;;
(* Declare i8* @gc_malloc(i64 size, i8* vistor_func ) *)
declare_function "gc_alloc" (function_type i8p_type [| i64_type; i8p_type |]) lmodule;;
(* Declare void @gc_add_root( i8** addr ) *)
declare_function "gc_add_root" (function_type void_type [| i8pp_type |]) lmodule;;

let codegen_gc_name rname = "gc_record_" ^ rname
;;

let codegen_gcroot addr rname =
  let addr = build_bitcast addr i8pp_type "rootptr" builder in
    ignore (codegen_call "llvm.gcroot" [| addr; const_null i8p_type |]);
;;

let codegen_malloc mptype =
  let visitor = match mptype with
    | Ast.Record rname -> build_bitcast (get_function (codegen_gc_name rname)) i8p_type "" builder
    | _ -> const_null i8p_type in
  let ltype = codegen_ltype mptype in
  let malloc = codegen_call "gc_alloc" [| (size_of @ element_type) ltype; visitor |] in
    (* Malloc returns i8* so we need to cast it to correct pointer type *)
    build_bitcast malloc ltype "gc_alloc" builder
;;

let rec codegen_addr_of = function
  | Ast.Var name -> Lenv.query name
  | Ast.Field (fname, expr, rname) ->
      let fptype = (pointer_type @ codegen_ltype) (Types.field_sort rname fname)
      and index = Types.field_index rname fname in
      let addr = build_struct_gep (codegen_expr expr) index fname builder in
        (* Field type could be i8* so we need to cast it back to something usable *)
        build_bitcast addr fptype "" builder 
  | _ -> failwith "can't get address of expression"
and codegen_expr = function 
  | Ast.T -> const_int boolean_type 1
  | Ast.F -> const_int boolean_type 0
  | Ast.Nil mptype -> const_null (codegen_ltype mptype)
  | Ast.Int value -> const_int integer_type value
  | Ast.Var name as var -> build_load (codegen_addr_of var) name builder
  | Ast.Field (fname, expr, rname) as field -> build_load (codegen_addr_of field) fname builder
  | Ast.Binexpr (bop, e1, e2) -> 
    begin match bop with
      | Ast.Plus -> build_add
      | Ast.Minus -> build_sub
      | Ast.Times -> build_mul
      | Ast.Div -> build_sdiv
      | Ast.Mod -> build_srem
      | Ast.Le -> build_icmp Icmp.Slt 
      | Ast.Leq -> build_icmp Icmp.Sle
      | Ast.Ge -> build_icmp Icmp.Sgt
      | Ast.Geq -> build_icmp Icmp.Sge
      | Ast.Eq -> build_icmp Icmp.Eq
      | Ast.Neq -> build_icmp Icmp.Ne
      | Ast.And -> build_and
      | Ast.Or -> build_or
    end (codegen_expr e1) (codegen_expr e2) "btmp" builder
  | Ast.Unexpr (Ast.Uminus, expr) -> build_sub (const_int integer_type 0) (codegen_expr expr) "uminus" builder
  | Ast.Unexpr (Ast.Neg, expr) -> build_sub (const_int boolean_type 1) (codegen_expr expr) "uneg" builder
  | Ast.Call (fname, exprs) ->
      let lvalue = codegen_call fname (Array.of_list (List.map codegen_expr exprs)) in
        match (Types.fun_sort fname) with
          | (_, (Ast.Record rname  as rettype)) ->
              (* If function returns pointer then mark it as gcroot *)
              let alloca = build_alloca (codegen_ltype rettype) fname builder in
                codegen_gcroot alloca rname;
                ignore (build_store lvalue alloca builder);
                build_load alloca "" builder
          | _ -> lvalue
and codegen_stat = function 
  | Ast.Assign (lhs, rhs) -> ignore (build_store (codegen_expr rhs) (codegen_addr_of lhs) builder)
  | Ast.New (expr, etype) -> ignore (build_store (codegen_malloc etype) (codegen_addr_of expr) builder)
  | Ast.If1 (expr, stat) -> codegen_stat (Ast.If2 (expr, stat, Ast.Seq []))
  | Ast.If2 (expr, tstat, fstat) ->
(* entry: br i1 %lexpr, label %then, label %else then: codegen_stat tstat  *)
(* br label %imerge else: codegen_stat fstat br label %imerge imerge:      *)
    let lexpr = codegen_expr expr in
    let start_bb = insertion_block builder in
    let lfunc = block_parent start_bb in
    let then_bb = append_block context "then" lfunc in
    position_at_end then_bb builder;
    codegen_stat tstat;
    let new_then_bb = insertion_block builder in
    let else_bb = append_block context "else" lfunc in
    position_at_end else_bb builder;
    codegen_stat fstat;
    let new_else_bb = insertion_block builder
    and merge_bb = append_block context "ifmerge" lfunc in
    position_at_end start_bb builder;
    ignore (build_cond_br lexpr then_bb else_bb builder);
    position_at_end new_then_bb builder; 
    ignore (build_br merge_bb builder);
    position_at_end new_else_bb builder;
    ignore (build_br merge_bb builder);
    position_at_end merge_bb builder
  | Ast.While (expr, stat) ->
(* entry: br label %test test: br i1 %lexpr, label %do, label %merge do:   *)
(* codegen_stat stat br label %test merge:                                 *)
    let entry_bb = insertion_block builder in
    let lfunc = block_parent entry_bb in
    let test_bb = append_block context "test" lfunc 
    and do_bb = append_block context "do" lfunc 
    and merge_bb = append_block context "merge" lfunc in
    position_at_end entry_bb builder;
    ignore (build_br test_bb builder);
    position_at_end test_bb builder;
    let lexpr = codegen_expr expr in
    ignore (build_cond_br lexpr do_bb merge_bb builder);
    position_at_end do_bb builder;
    codegen_stat stat;
    ignore (build_br test_bb builder);
    position_at_end merge_bb builder;
  | Ast.Seq statements -> List.iter codegen_stat statements
  | Ast.Expr expr -> ignore(codegen_expr expr)

let codegen_alloca builder (vname, vtype) =
  let alloca = build_alloca (codegen_ltype vtype) vname builder in
    Lenv.push vname alloca;
    match vtype with
      | Ast.Record rname -> codegen_gcroot alloca rname; alloca
      | _ -> alloca
;;

let codegen_declare fname (args, ret) =
  let largstype = Array.of_list (List.map codegen_ltype args) in
  let lftype = function_type (codegen_ltype ret) largstype in
  let lfun = declare_function fname lftype lmodule in
    Lenv.push fname lfun;
    lfun

let codegen_func (func : Ast.mpfun_t) =
  let lfun = codegen_declare func#name (List.map snd func#parms, func#rettype) in
  let bb = append_block context "entry" lfun in
    set_gc (Some "shadow-stack") lfun;
    position_at_end bb builder;
    (* Generate alloca for arguments so we can modify them *)
    Array.iteri (fun i value ->
      let param = List.nth func#parms i in
      (* Set name for function argument *)
      set_value_name (fst param) value;
      (* Allocate memory on stack for argument *)
      let alloca = codegen_alloca builder param in
      (* Store value from argument in newly created alloca *)
      ignore (build_store value alloca builder)
    ) (params lfun);
    (* Generate alloca for local variables *)
    List.iter (ignore @ codegen_alloca builder) func#locals; 
    (* Create alloca for return value *)
    if func#rettype <> Ast.Void then
      ignore (codegen_alloca builder (func#name, func#rettype));
    (* Generate function body *)
    codegen_stat func#body;
    (* Return value *)
    if func#rettype <> Ast.Void then begin
      let lret_value = build_load (Lenv.query func#name) func#name builder in
        ignore (build_ret lret_value builder); 
        Lenv.pop func#name;
    end else
      ignore (build_ret_void builder);
    Llvm_analysis.assert_valid_function lfun;
    (* Remove bindings (locals and params) *)
    List.iter (Lenv.pop @ fst) func#locals;
    List.iter Lenv.pop (List.map fst func#parms);
    ()

let codegen_global (vname, vtype) =
  let ltype = codegen_ltype vtype in
    Lenv.push vname (define_global vname (const_null ltype) lmodule)

let codegen_gc_object_type =
  (* struct gc_object { struct gc_object *next; unsigned char marked;  gc_visitor visitor; unsigned int size }; *)
  pointer_type (struct_type context [| i8p_type; boolean_type; i8p_type; integer_type |])
;;

let codegen_visit rname =
  let lfun = declare_function (codegen_gc_name rname) (function_type void_type [| codegen_gc_object_type |] ) lmodule in
  let entry = append_block context "entry" lfun in
    position_at_end entry builder;
    let gcparam = param lfun 0 in
    let notnullbl = append_block context "if_not_null" lfun in
    let notmarkedbl = append_block context "if_not_marked" lfun in
    let mergebl = append_block context "merge_block" lfun in
      position_at_end (entry_block lfun) builder;
      ignore (build_cond_br (build_icmp Icmp.Ne (const_null codegen_gc_object_type) gcparam "" builder) notnullbl mergebl builder);
      position_at_end notnullbl builder;
      let markedaddr = build_gep gcparam [| const_int i32_type (- 1); const_int i32_type 1 |] "" builder in
        ignore (build_cond_br (build_load markedaddr "" builder) mergebl notmarkedbl builder);
        position_at_end notmarkedbl builder;
        ignore (build_store (const_int boolean_type 1) markedaddr builder);
        let fields = Types.fields_sort rname in
	  List.iter (
	    function
	      | (fname, Ast.Record frname) ->
	          let record = build_bitcast gcparam (codegen_ltype (Ast.Record rname)) "" builder in
	          let field_addr_i8p = build_struct_gep record (Types.field_index rname fname) "" builder in
                  let field_addr_i8 = build_load field_addr_i8p "" builder in
                  let field_addr = build_bitcast field_addr_i8 codegen_gc_object_type "" builder in
	            ignore (codegen_call "gc_visit" [| field_addr |] )
	      | _ -> ()
	  ) fields;
          ignore (build_br mergebl builder);
	  position_at_end mergebl builder;
          ignore (build_ret_void builder)
            
let codegen_prog (prog : Ast.mpprog_t) =
  List.iter codegen_global prog#vars;
  let gc_visit_func_type = function_type void_type [| codegen_gc_object_type |] in
  ignore (declare_function "gc_visit" gc_visit_func_type lmodule);
  List.iter (fun record -> ignore (declare_function (codegen_gc_name record#name) gc_visit_func_type lmodule)) prog#recorddecls;
  List.iter (fun record -> codegen_visit record#name) prog#recorddecls;
  List.iter (fun (fname, ftype) -> ignore(codegen_declare fname ftype)) prog#fundecls;
  List.iter codegen_func prog#funs;
  codegen_func (Ast.new_mpfun "main" [] [] prog#body Ast.Integer);
  let entry = get_function "main" in
    position_builder (instr_begin (entry_block entry)) builder;
    List.iter (
      fun (vname, vtype) -> match vtype with
        | Ast.Record rname ->
            let addr = build_bitcast (Lenv.query vname) i8pp_type "addr" builder in
              ignore (codegen_call "gc_add_root" [| addr |])
        | _ -> ()
    ) prog#vars

let name = ref ""
;;

let set_name s =
  name := s
;;

let the_fpm = PassManager.create ()
;;
add_cfg_simplification the_fpm;
add_memory_to_register_promotion the_fpm;
add_instruction_combination the_fpm;
add_reassociation the_fpm;
add_gvn the_fpm;
add_instruction_combination the_fpm;
add_cfg_simplification the_fpm;

Arg.parse [] set_name "codegen input";
Run.run (fun name -> ignore (codegen_prog (Parsemp.parse name))) !name;
Llvm_analysis.assert_valid_module lmodule;
ignore (PassManager.run_module lmodule the_fpm);
Llvm_analysis.assert_valid_module lmodule;
dump_module lmodule;
ignore(Llvm_bitwriter.write_bitcode_file lmodule (!name ^ ".bc"));
dispose_module lmodule
